#include <limits>
#include "helper/helper.h"
#include "helper/arma_helper.h"

// make negative elements zeros
void zero_threshold(arma::mat& A)
{
  for (unsigned int i = 0; i < A.n_rows; i++)
  for (unsigned int j = 0; j < A.n_cols; j++)
  if (A(i,j) < 1e-16) A(i,j) = 0;
}

// make column to norm 1 (projection on sphere)
void normalize_column(arma::mat& A)
{
  for (unsigned int j = 0; j < A.n_cols; j++) {
    double s = arma::norm(A.col(j), 2);
    if (s > 1e-16)
    A.col(j) /= s;
  }
}

// soft thresholding
void soft_threshold(arma::mat& A, double k) {
  for (unsigned int i = 0; i < A.n_rows; i++)
    for (unsigned int j = 0; j < A.n_cols; j++) {
      double v = A(i,j);
      if (v > k) A(i,j) = v-k;
      else if (v < -k) A(i,j) = v+k;
      else A(i,j) = 0;
    }
}

void normalize_column_less_1(arma::mat& V)
{
  for (unsigned int j = 0; j < V.n_cols; j++) {
    double s = norm(V.col(j), 2);
    if (s > 1) V.col(j) /= s;
  }
}

void normalize_column_equal_1(arma::mat& V)
{
  for (unsigned int j = 0; j < V.n_cols; j++) {
    // scan for all positive element
    double s = 0;
    int zero_index = -1;
    double max_negative = -std::numeric_limits<double>::infinity();
    int max_index = -1;
    for (unsigned int i = 0; i < V.n_rows; i++) {
      double v = V(i,j);
      if (v >= 1e-16) s += v*v;
      else if (fabs(v) < 1e-16) zero_index = i;
      else if (v > max_negative) {
        max_negative = v;
        max_index = i;
      }
    }

    if (s > 0) { // there is positive element
      for (unsigned int i = 0; i < V.n_rows; i++)
        if (V(i,j) < 1e-16) V(i,j) = 0;
      V.col(j) /= sqrt(s);
    }
    else { // there is no positive element
      int one_index = (zero_index != -1) ? zero_index : max_index;
      if (one_index == -1) FATAL("SOMETHING WRONG WITH PROJECTION");
      for (unsigned int i = 0; i < V.n_rows; i++)
        if ((int) i == one_index) V(i,j) = 1;
        else V(i,j) = 0;
    }
  }
}
